import os
import random
import shutil
from collections import defaultdict

# Set random seed for reproducibility
random.seed(42)

# Paths to the data directories
data_dirs = [
    'd:/competition/2026quanqiu/data/webfg400_train/train',
    'd:/competition/2026quanqiu/data/webinat5000_train/train'
]

# Directory to store the split datasets
output_dir = 'd:/competition/2026quanqiu/split_data'

# Function to split dataset
def split_dataset(data_dir, output_dir, dataset_name, train_ratio=0.8):
    print(f"Processing dataset in {data_dir}")
    # Organize images with labels
    for label in os.listdir(data_dir):
        label_dir = os.path.join(data_dir, label)
        # Check if label directory exists
        if not os.path.exists(label_dir):
            print(f"Directory does not exist: {label_dir}")
            continue

        # List images in the label directory
        images = []
        for img in os.listdir(label_dir):
            image_path = os.path.join(label_dir, img)
            if os.path.isfile(image_path):
                images.append(image_path)

        random.shuffle(images)
        print(f"Total images in class '{label}': {len(images)}")

        # Split files into training and validation
        split_index = int(len(images) * train_ratio)
        train_images = images[:split_index]
        val_images = images[split_index:]
        print(f"Class '{label}' - Training images: {len(train_images)}, Validation images: {len(val_images)}")

        # Create directories for train and val
        train_label_dir = os.path.join(output_dir, dataset_name + '_train', label)
        val_label_dir = os.path.join(output_dir, dataset_name + '_val', label)
        os.makedirs(train_label_dir, exist_ok=True)
        os.makedirs(val_label_dir, exist_ok=True)

        # Copy files to respective directories
        for image_path in train_images:
            label_dir = os.path.join(train_label_dir, os.path.basename(image_path))
            print(f"Copying {image_path} to {label_dir}")
            shutil.copy(image_path, train_label_dir)
        for image_path in val_images:
            label_dir = os.path.join(val_label_dir, os.path.basename(image_path))
            print(f"Copying {image_path} to {label_dir}")
            shutil.copy(image_path, val_label_dir)

    print(f"Finished processing {data_dir}")

# Process each data directory
for data_dir in data_dirs:
    dataset_name = os.path.basename(os.path.dirname(data_dir))
    split_dataset(data_dir, output_dir, dataset_name)
